#include "znx_pool.h"


static znx_pool_page_t *
znx_pool_page_alloc(size_t page_size)
{
    if (page_size == 0) {
        page_size = ZNX_POOL_PAGE_SIZE;
    }

    if (page_size < ZNX_POOL_MIN_PAGE_SIZE) {
        page_size = ZNX_POOL_MIN_PAGE_SIZE;
    }

    znx_pool_page_t *page = znx_thread_calloc(page_size);
    page->start = (u_char *)page + sizeof(znx_pool_page_t);
    page->pos = page->start;
    page->end = (u_char *)page + page_size;
    return page;
}


static void
znx_pool_page_free(znx_pool_page_t *page)
{
    znx_thread_free(page);
}


static int
znx_pool_block_node_cmp(znx_rbtree_node_t *a, znx_rbtree_node_t *b)
{
    znx_pool_block_node_t *a_node, *b_node;
    a_node = (znx_pool_block_node_t *)a;
    b_node = (znx_pool_block_node_t *)b;

    if (a_node->ptr == b_node->ptr) {
        return 0;
    }

    return a_node->ptr > b_node->ptr ? 1 : -1;
}


void *
znx_pool_block_malloc(znx_pool_t *pool, size_t size)
{
    void *ptr = znx_thread_malloc(size);
    znx_pool_block_node_t *block_node
        = znx_thread_malloc(sizeof(znx_pool_block_node_t));
    block_node->ptr = ptr;
    znx_rbtree_insert(&pool->block_tree, &block_node->tree_node);
    return ptr;
}


void
znx_pool_block_free(znx_pool_t *pool, void *ptr)
{
    znx_pool_block_node_t   target = {.ptr = ptr};
    znx_rbtree_node_t       *node;
    znx_pool_block_node_t   *block_node;

    node = znx_rbtree_lookup(&pool->block_tree, &target.tree_node);
    if (!node) {
        return; // not found
    }

    znx_rbtree_delete(&pool->block_tree, node);
    block_node = (znx_pool_block_node_t *)node;

    znx_thread_free(block_node->ptr);
    znx_thread_free(block_node);
}


znx_pool_t *
znx_pool_create()
{
    znx_pool_t *pool = znx_thread_calloc(sizeof(znx_pool_t));
    pool->page_size = ZNX_POOL_MIN_PAGE_SIZE;
    znx_queue_init(&pool->pages);
    znx_rbtree_init(&pool->block_tree,
        &pool->block_sentinel, znx_pool_block_node_cmp);
    znx_cleanup_list_init(&pool->cleanup_list);
    return pool;
}


znx_pool_t *
znx_pool_create_ext(size_t page_size)
{
    znx_pool_t *pool = znx_thread_calloc(sizeof(znx_pool_t));
    pool->page_size = page_size;
    znx_queue_init(&pool->pages);
    znx_rbtree_init(&pool->block_tree,
        &pool->block_sentinel, znx_pool_block_node_cmp);
    znx_cleanup_list_init(&pool->cleanup_list);
    return pool;
}


void
znx_pool_destroy(znx_pool_t *pool)
{
    znx_pool_page_t         *page;
    znx_rbtree_node_t       *node;
    znx_pool_block_node_t   *block_node;
    znx_queue_t             *q;

    znx_cleanup_list_exec(&pool->cleanup_list);

    while (!znx_rbtree_empty(&pool->block_tree)) {
        node = znx_rbtree_min(pool->block_tree.root, pool->block_tree.sentinel);
        znx_rbtree_delete(&pool->block_tree, node);
        block_node = (znx_pool_block_node_t *)node;
        if (block_node->ptr != NULL) {
            znx_thread_free(block_node->ptr);
        }
        znx_thread_free(block_node);
    }

    while (!znx_queue_empty(&pool->pages)) {
        q = znx_queue_last(&pool->pages);
        znx_queue_remove(q);
        page = znx_queue_data(q, znx_pool_page_t, queue);
        znx_pool_page_free(page);
    }

    znx_thread_free(pool);
}


void *
znx_pool_malloc(znx_pool_t *pool, size_t size)
{
    size = (size_t)ZNX_ALIGN(size, (size_t)ZNX_POOL_ALLOC_ALIGNED_SIZE);

#if (ZNX_POOL_NO_PAGE)
    return znx_pool_block_malloc(pool, size);
#endif

    if (size > ZNX_POOL_MIN_PAGE_SIZE) {
        return znx_pool_block_malloc(pool, size);
    }

    znx_pool_page_t *page = pool->cur_page;
    if (page == NULL || (size_t)(page->end - page->pos) < size) {
        page = znx_pool_page_alloc(pool->page_size);
        znx_queue_insert_head(&pool->pages, &page->queue);
        pool->cur_page = page;
    }

    void *ptr = page->pos;
    page->pos += size;

    return ptr;
}


void
znx_pool_free(znx_pool_t *pool, void *ptr)
{
    // relase block mem only
    znx_pool_block_free(pool, ptr);
}
